import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import csv


# 读取3d的.csv
def read_from_csv(path_csv):
    x_coords = []
    y_coords = []
    z_values = []

    with open(path_csv, 'r') as csvfile:
        reader = csv.DictReader(csvfile)
        for row in reader:
            x_coords.append(float(row['X']))
            y_coords.append(float(row['Y']))
            z_values.append(float(row['Z']))

    x_coords = np.unique(np.array(x_coords))
    y_coords = np.unique(np.array(y_coords))
    z_values = np.array(z_values).reshape((len(x_coords), len(y_coords)))

    return x_coords, y_coords, z_values


# HFSS方向图转成360*360矩阵
def format_hfss3d_2_360x360(path_hfss3d, path_360x360):
    pattern = np.zeros((360, 90))
    pattern_line = np.array(pd.read_csv(path_hfss3d))

    for ii in pattern_line:
        if ii[0] >= 0:
            if ii[0] < 360 and ii[1] < 90:
                pattern[int(ii[0])][int(ii[1])] = ii[2]
        else:
            pattern[int(360 + ii[0])][int(ii[1])] = ii[2]

    pattern = pd.DataFrame(pattern)
    pattern.to_csv(path_360x360, index=False)


def plot_pattern3d(path_360x360):
    pattern = np.array(pd.read_csv(path_360x360))
    phi = np.arange(0, 359, 1)  # 0 to 360 degrees
    theta = np.arange(0, 90, 1)  # 0 to 180 degrees
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')
    Theta, Phi = np.meshgrid(theta, phi)
    pattern = pattern / np.max(pattern)
    ex = (pattern * np.sin(Theta)) * np.cos(Phi)
    ey = (pattern * np.sin(Theta)) * np.sin(Phi)
    ez = pattern * np.cos(Theta)
    ax.plot_surface(ex, ey, ez, rstride=1, cstride=1, cmap=plt.cm.coolwarm)
    ax.set_xlim3d(-1, 1)
    ax.set_ylim3d(-1, 1)
    ax.set_zlim3d(0, 1)
    ax.set_xlabel('x')
    ax.set_ylabel('y')
    plt.show()

def plot_pattern3d_2(path_hfss3d):
    x_coords, y_coords, z_matrix = read_from_csv(path_hfss3d)
    #
    fig = plt.figure(figsize=(10, 8))
    ax = fig.add_subplot(111, projection='3d')
    X, Y = np.meshgrid(x_coords, y_coords)
    Z = z_matrix.T
    ax.scatter(X, Y, Z, c=Z, cmap='gist_rainbow')
    # 隐藏轴标签
    ax.set_xlabel('')
    ax.set_ylabel('')
    ax.set_zlabel('')
    # 显示网格线
    ax.grid(True)
    plt.show()


def plot_pattern3d_3(path_hfss3d):
    """
    绘制3D方向图，正确地将(theta, phi, r)转换为(x, y, z)。
    """
    # 1. 读取数据
    theta_coords, phi_coords, r_matrix = read_from_csv(path_hfss3d)

    # 2. 创建Theta和Phi的网格 (弧度)
    Theta_grid, Phi_grid = np.meshgrid(np.radians(theta_coords), np.radians(phi_coords), indexing='ij')
    R_grid = r_matrix

    # 3. 球坐标转笛卡尔坐标 (假设Theta是从Z轴正方向测量的俯仰角)
    X = np.sin(Theta_grid) * np.cos(Phi_grid)
    Y = np.sin(Theta_grid) * np.sin(Phi_grid)
    Z = np.cos(Theta_grid)

    # 4. 绘图
    fig = plt.figure(figsize=(12, 8))
    ax = fig.add_subplot(111, projection='3d')

    # 使用转换后的(X, Y, Z)进行绘图
    # 可以选择绘制表面图或散点图
    # --- 表面图 ---
    # surf = ax.plot_surface(X, Y, Z, facecolors=plt.cm.gist_rainbow(R_grid/R_grid.max()), alpha=0.8)
    # m = plt.cm.ScalarMappable(cmap=plt.cm.gist_rainbow)
    # m.set_array(R_grid)
    # fig.colorbar(m, ax=ax, shrink=0.5, aspect=5, label='Field Strength (mV)')

    # --- 散点图 ---
    # 将网格数据展平以便散点图绘制
    sc = ax.scatter(X.ravel(), Y.ravel(), Z.ravel(), c=R_grid.ravel(), cmap='gist_rainbow', s=1)
    fig.colorbar(sc, ax=ax, shrink=0.5, aspect=5, label='Field Strength (mV)')

    # 设置图表属性
    ax.set_title('3D Radiation Pattern')
    # 隐藏坐标轴标签，因为它们是X,Y,Z笛卡尔坐标，而非原始的Theta,Phi
    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.set_zlabel('Z')
    ax.grid(True)

    # 设置坐标轴比例相等
    # max_range = np.array([X.ravel().max() - X.ravel().min(), Y.ravel().max() - Y.ravel().min(),
    #                       Z.ravel().max() - Z.ravel().min()]).max() / 2.0
    # mid_x = (X.ravel().max() + X.ravel().min()) * 0.5
    # mid_y = (Y.ravel().max() + Y.ravel().min()) * 0.5
    # mid_z = (Z.ravel().max() + Z.ravel().min()) * 0.5
    # ax.set_xlim(mid_x - max_range, mid_x + max_range)
    # ax.set_ylim(mid_y - max_range, mid_y + max_range)
    # ax.set_zlim(mid_z - max_range, mid_z + max_range)

    plt.tight_layout()
    plt.show()




if __name__=="__main__":
    # 转HFSS输出为360x360
    # format_hfss3d_2_360x360("../files/tools/plot/32x32-beam1(30,0)-rE_Plot_3d.csv",
    #                         "../files/tools/plot/32x32-beam1(30,0)-rE_Plot_3d_360x360.csv")
    # 画图
    plot_pattern3d("../files/tools/plot/32x32-beam1(0,0)-rE_Plot_3d_360x360.csv")
    plot_pattern3d_2("../files/tools/plot/32x32-beam1(0,0)-rE_Plot_3d_2.csv")
    plot_pattern3d_3("../files/tools/plot/32x32-beam1(0,0)-rE_Plot_3d_2.csv")